import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl

class DiagnosisModel(pl.LightningModule):
    def __init__(self, input_dim=1024, num_classes=5, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.lr = lr
        
        self.conv_layers = nn.Sequential(
            nn.Conv1d(1, 16, 33, padding=16), # [B, 16, 4096]
            nn.BatchNorm1d(16), 
            nn.ReLU(),
            nn.MaxPool1d(32), # [B, 16, 128]
            
            nn.Conv1d(16, 32, 9, padding=4), # [B, 32, 16]
            nn.BatchNorm1d(32), 
            nn.ReLU(), 
            nn.MaxPool1d(8), # [B, 32, 2]
            
            nn.Conv1d(32, 64, 3, padding=1), # [B, 64, 16]
            nn.BatchNorm1d(64), 
            nn.ReLU(), 
            nn.AdaptiveAvgPool1d(16) # [B, 64, 16]
        )
        
        self.classifier = nn.Sequential(
            nn.Linear(64*16, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        # x: [B, 1, 4096]
        features = self.conv_layers(x) # [B, 64, 16]
        return self.classifier(features.view(x.size(0), -1)), features # [B, 5], [B, 64, 16]

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits, _ = self(x)
        loss = F.cross_entropy(logits, y)
        
        preds = torch.argmax(logits, dim=1)
        acc = (preds == y).float().mean()
        
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_acc', acc, prog_bar=True)
        
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits, _ = self(x)
        loss = F.cross_entropy(logits, y)
        
        preds = torch.argmax(logits, dim=1)
        acc = (preds == y).float().mean()
        
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, 
            mode='min', 
            factor=0.5, 
            patience=5, 
            verbose=True
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss",
                "interval": "epoch",
                "frequency": 1
            }
        }

    def get_condition_features(self, x):
        """获取条件特征
        Args:
            x: 输入张量 [B, 1, 1024]
        """
        # 确保输入在正确的设备上
        device = next(self.parameters()).device
        if x.device != device:
            print(f"Warning: Input device ({x.device}) does not match model device ({device}). Moving input to {device}")
            x = x.to(device)
        
        x = self.conv_layers[0](x)
        x = self.conv_layers[1:4](x)
        x = self.conv_layers[4:](x)
        
        return x

